#!/usr/bin/env python3

'''
    Software License

    Copyright (C) 2021-05-24  Xoronos

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, version 3.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''

'''
    Liabilities

    The software is provided "AS IS" without any warranty of any kind, either expressed,
    implied, or statutory, including, but not limited to, any warranty that the software
    will conform to specifications, any implied warranties of merchantability, fitness
    for a particular purpose, and freedom from infringement, and any warranty that the
    documentation will conform to the software, or any warranty that the software will
    be error free.

    In no event shall Xoronos be liable for any damages, including, but not limited to,
    direct, indirect, special or consequential damages, arising out of, resulting from,
    or in any way connected with this software, whether or not based upon warranty,
    contract, tort, or otherwise, whether or not injury was sustained by persons or
    property or otherwise, and whether or not loss was sustained from, or arose out of
    the results of, or use of, the software or services provided hereunder.
    
    To request the provided software under a different license you can contact us at
    support@xoronos.com
'''

import argparse
import random 
import re

import requests
import json
import os

from PIL import Image, ImageDraw
from cmath import phase
from sys import maxsize
from math import log10

folder = "./gen/"
size = "256x256"
localai_server_image = "http://localhost:8080/v1/images/generations"
localai_server_chat = "http://localhost:8080/v1/chat/completions"
localai_server_t2s = 'http://localhost:8080/tts'

# This is the default implementation with localai 
# for implementing a different service you can reimplement gen, get_names, flush ,delete_all 
class XRN_gen_text:
    #def __init__(self):

    def gen(self,prompt):

        self.headers = {
            "Content-Type": "application/json"
        }

        self.json_data = {
            "model": "gpt-4",
            "messages": [{"role": "user", "content": prompt }] 
        }

        response = requests.post(localai_server_chat, headers=self.headers, data=json.dumps(self.json_data))
        
        if response.status_code == 200 :
            
            t = json.loads(response.text)

            for choiche in  t['choices'] :
                return choiche["message"]["content"]
        else:
            print("responce not valid ")
            print(response.text)  # Print the content of the response
            return -1
        return 0

# This is the default implementation with localai 
# for implementing a different service you can reimplement gen, get_names, flush ,delete_all 
class XRN_gen_image:
    def __init__(self):
        
        self.files= []

    def gen(self,prompt,fwpath=""):

        self.headers = {
            "Content-Type": "application/json"
        }

        self.json_data = {
            "prompt": prompt,
            "size": size
        }
        response = requests.post(localai_server_image, headers=self.headers, data=json.dumps(self.json_data))
        
        if response.status_code == 200 :
            
            t = json.loads(response.text)

            for a in t['data'] :
                print (a['url'])
                image_url = a['url']

            splitteturl = image_url.split("/")

            for x in splitteturl:
                lastname = x

            #print(lastname)
            self.files.append(lastname)

            try:
                response = requests.get(image_url)
                response.raise_for_status()  # Check if the request was successful
                if fwpath == "":
                    fwpath = folder+lastname
                with open(fwpath, 'wb') as file:
                    file.write(response.content)
            except Exception as err:
                print(f' error occurred: {err}')  # Print any other error
                return -2
        else:
            print("responce not valid ")
            print(response.text)  # Print the content of the response
            return False
        return True
    def get_names(self):
        return self.files
    def flush(self):
        self.files = []
    def delete_all(self):
        for n in self.files:
            os.remove(folder+n)

class XRN_gen_audio:
    def __init__(self):
        
        self.files= []
        random.seed(version=2)
            
        self.name = random.randbytes(8).hex()
        self.audiocount = 0 

    def gen(self,prompt,filename=None):

        self.headers = {
            "Content-Type": "application/json"
        }

        self.json_data = {
            "backend": "auto",
            "model": "tts-1",
            "input": prompt

        }
        try:
            response = requests.post(localai_server_t2s, headers=self.headers, data=json.dumps(self.json_data))
            response.raise_for_status()  # Check if the request was successful
            if filename == None :
                filename = folder+self.name+str(self.audiocount)+".wav"
            with open( filename, 'wb') as file:
                file.write(response.content)
            self.files.append(filename)
            self.audiocount += 1

            
        except Exception as err:
            print(f'Other error occurred: {err}')  # Print any other error
            return False
        return True
    def get_names(self):
        return self.files
    def flush(self):
        self.files = []
    def delete_all(self):
        for n in self.files:
            os.remove(folder+n)

def RGBtoHSV(R, G, B):
    R /= 255 ; G /= 255 ; B /= 255
    cmin = min(R, G, B); cmax = max(R, G, B)
    dmax = cmax - cmin ;  V = cmax
    
    if dmax == 0:
        H = 0 ; S = 0
    
    else:
        S = dmax/cmax
    
        dR = ((cmax - R)/6 + dmax/2)/dmax
        dG = ((cmax - G)/6 + dmax/2)/dmax
        dB = ((cmax - B)/6 + dmax/2)/dmax
    
        if   R == cmax: H = (dB - dG)%1
        elif G == cmax: H = (1/3 + dR - dB)%1
        elif B == cmax: H = (2/3 + dG - dR)%1
    
    return (H, S, V)

def generate_fractal(sizex, sizey, image_path) :
    dim  = (sizex, sizey)
    bits = 8
    cmax = (1<<bits)-1
    cfac = 255/cmax
    
    img  = Image.new('RGB', dim)
    draw = ImageDraw.Draw(img)
    
    start = random.randrange( -100 , 100 ) / random.randrange( 1 , 200 )
    focal = random.randrange(1, 1000) / random.randrange( 1, 1000)
    xstart = start ; ystart = start
    
    
    xd = focal / dim[0] ; yd = focal / dim[1]
    
    tol = 1e-6
    a = [[], [], [], [], []]
    
    for x in range(dim[0]):
        # print(x, "\r", end=' ')
        # print(x, " ")
        for y in range(dim[1]):
            z = d = complex(xstart + x*xd, ystart + y*yd)
            c = 0.0
            l = 1
            while abs(l-z) > tol and abs(z) > tol:
                l = z
                z -= (z**5-1)/(5*z**4)
                c += 1.0
            if z == 0: 
                c = maxsize
            p = int(phase(z))
            
            if abs(l-z) > 0.0:
                c += log10(abs(l-z)) / 6
            
            a[p] += (c, x, y),
    
    for i in range(5):
        a[i].sort(reverse = False)
    
    pnum = [len(a[i]) for i in range(5)]
    ptot = dim[0]*dim[1]
    
    bounds = []
    lbound = 0
    for i in range(4):
        nbound = lbound + pnum[i]/ptot
        bounds += nbound,
        lbound = nbound

    t = [[], [], [], [], []]
    for i in range(ptot-1, -1, -1):
        r = (i>>bits*2)*cfac
        g = (cmax&i>>bits)*cfac
        b = (cmax&i)*cfac
        (h, s, v) = RGBtoHSV(r, g, b)
        h = (h+0.1)%1
        if   h < bounds[0] and len(t[0]) < pnum[0]: p=0
        elif h < bounds[1] and len(t[1]) < pnum[1]: p=1
        elif h < bounds[2] and len(t[2]) < pnum[2]: p=2
        elif h < bounds[3] and len(t[3]) < pnum[3]: p=3
        else: p=4
        t[p] += (int(r), int(g), int(b)),


    rt = [0, 0, 0, 0, 0]
    for i in range( 5 ):
        t[i].sort(key = lambda c: c[0]*2126 + c[1]*7152 + c[2]*722, reverse = True)
        rt[i] = len(t[i])

    r = [0, 0, 0, 0, 0]
    rc = 0 ; rb = 0 ; rg = 0;

    imgmatrix = [[(0,0,0) for _ in range(sizey)] for _ in range(sizex)]

    for p in range( 5 ):
        for c,x,y in a[p]:
            if ( r[p] < rt[p] ) :
                imgmatrix[x][y] = t[p][r[p]]
                r[p] += 1

    shadex=0; shadey=0; rndshade=random.randint(0,7)
    for x in range (sizex) :
        if ( x % (sizex / 256) == 0 ) : shadex = shadex + 1 % 256
        for y in range (sizey) :
            if ( y % (sizey / 256) == 0 ) : shadey = shadey + 1 % 256
            if (0,0,0) == imgmatrix[x][y] :
                if   rndshade ==  0 : imgmatrix[x][y] = (0, shadex, 0)
                elif rndshade ==  1 : imgmatrix[x][y] = (0, shadey, 0)
                elif rndshade ==  2 : imgmatrix[x][y] = (0, 0, shadex)
                elif rndshade ==  3 : imgmatrix[x][y] = (0, 0, shadey)
                elif rndshade ==  4 : imgmatrix[x][y] = (0, shadex, shadex)
                elif rndshade ==  5 : imgmatrix[x][y] = (0, shadey, shadey)
                elif rndshade ==  6 : imgmatrix[x][y] = (0, shadex, shadey)
                elif rndshade ==  7 : imgmatrix[x][y] = (0, shadey, shadex)

            if random.randint(0,4) == 0 : imgmatrix[x][y] = tuple(map(sum,zip(imgmatrix[x][y], (random.randint(0,1),random.randint(0,1),random.randint(0,1) ) )))
            draw.point((x,y), imgmatrix[x][y])

    img.save(image_path)


concat_N = 4
def main():
    parser = argparse.ArgumentParser(
        description="Utility for generating content",
        usage="""%(prog)s [--png [-prompt "string" | -auto-prompt | -no-prompt ] -size WITHxHEIGH -image /path/to/image.png ] [--flac [ -prompt "string" | -auto-prompt ] -audio /path/to/audio.flac ]
        
Examples:

"""
    )

    # Adding arguments for --flac
    parser.add_argument('--flac', action='store_true', help="Text to speech  command")
    parser.add_argument('-audio', type=str, help="First argument for flac (mandatory if --flac is used)")

    
    # Adding arguments for --png
    parser.add_argument('--png', action='store_true', help="Text to image command")
    parser.add_argument('-size', type=str,nargs=1,help="Size of image WITHxHEIGH")
    parser.add_argument('-image', type=str,nargs=1, help="Third argument for png (mandatory if --png is used)")
    
    # Adding mutually exclusive arguments
    exclusive_group = parser.add_mutually_exclusive_group()
    exclusive_group.add_argument('-prompt', help='the prompt to send ',type=str,nargs=1)
    exclusive_group.add_argument('-auto-prompt', action='store_true', help='autogenerate prompt')
    exclusive_group.add_argument('-no-prompt', action='store_true', help='generate a fractal image')
    
    # Parsing arguments
    args = parser.parse_args()
    
    # Handling arguments
    if args.flac:
        if not check_prompt(args) :
            parser.print_help()
        else : 
            if args.audio:
                handle_flac(args)
            else:
                print("Error: --flac requires -audio.")
                parser.print_usage()
    elif args.png:
        pattern = r'^\d+(x|X)\d+$'
        
        if not check_prompt(args) :
            parser.print_help()
        else : 
            if args.size and args.image :
                match = bool(re.match(pattern, args.size[0])) 
                if match : 
                    handle_png(args)
                else:
                    print("Error: -size argument sintax wrong ")
                    parser.print_usage()


            else:
                print("Error: -image requires -size and -image")
                parser.print_usage()
    else:
        parser.print_help()

def check_prompt(args):
    if args.prompt or args.auto_prompt or args.no_prompt :
        if ( args.prompt and args.auto_prompt ) or ( args.no_prompt and args.auto_prompt ) or ( args.no_prompt and args.prompt ) :
            print("Error: only one between -prompt, -auto-prompt or -no-prompt can be used")
            return False
        return True
    else :
        print("Error: -prompt or -auto-prompt or -no-prompt is required")
        return False
    
def gen_auto_prompt():
    print("auto_prompt could be very slow, if you want to speedup it is possible to use -prompt \"Prompt sentence \" ")
    auto_prompt_list = []
    with open("code/autoprompt.txt","r") as file:
        line_count = 0
        for line in file:
            #print(line)
            # Increment the line count
            line_count += 1
            # Print the current line
            auto_prompt_list.append(line)
    #print(line_count)
    auto_prompt = ""
    for i in range(0,concat_N) :
        auto_prompt += auto_prompt_list[random.randint(0, line_count-1)]
    auto_prompt = auto_prompt.replace('\n', ' ')
    #print(auto_prompt)

    x = XRN_gen_text()
    return  x.gen(auto_prompt)


def handle_flac(args):
    if args.auto_prompt:
        prompt = gen_auto_prompt() 
    else :
        prompt = args.prompt[0]
    print(prompt)        
    x = XRN_gen_audio()
    x.gen(prompt)

def handle_png(args):
    if args.auto_prompt:
        prompt = gen_auto_prompt() 
        x = XRN_gen_image()
        x.gen(prompt,args.image[0])
    elif args.no_prompt:
        if ( 'X' in args.size[0] ) :
            x, y = args.size[0].split('X')
        elif ( 'x' in args.size[0] ) :
            x, y = args.size[0].split('x')
        x = int (x)
        y = int (y)
        generate_fractal(x, y, args.image[0])
    else :
        prompt = args.prompt[0]
        x = XRN_gen_image()
        x.gen(prompt,args.image[0])

    
if __name__ == "__main__":
    main()
